# -*- encoding: utf-8 -*-

"""
@File:      disk.py
@Time:      2022/08/15 16:18:47
@Author:    zhangtaibo.ztb
@Version:   1.0
@Contact:   zhangtaibo.ztb@alibaba-inc.com
@License:   Mulan PSL v2
"""

from common.sysinfo import SysInfo

class DiskManager():
    """
    Common def for disk manager.
    """

    def __init__(self):
        self.cmd = SysInfo.cmd()

    def prepare_datadisk(self, host="local", test_img="/tmp/test100.img", mnt_dir="/mnt/testpoint", fs_type="ext4", fs_opt=""):
        self.cmd("dd if=/dev/zero of=%s bs=1M count=100" % test_img, host)
        self.cmd("mkfs.%s %s %s" % (fs_type, fs_opt, test_img), host)
        self.cmd("ls %s && rm -rf %s" % (mnt_dir, mnt_dir), host, ignore_status=True)
        self.cmd("mkdir -p %s" % mnt_dir, host)
        self.cmd("mount %s %s" % (test_img, mnt_dir), host)
        ret_c, data_disk = self.cmd('df -TH |grep -w "%s .* %s" |awk "{print \$1}"' % (fs_type, mnt_dir), host)
        return data_disk

    def cleanup_datadisk(self, test_img, mnt_dir, host="local"):
        self.cmd("umount %s" % mnt_dir, host, ignore_status=True)
        self.cmd("rm -rf %s %s" % (mnt_dir, test_img), host)

    def get_data_disk(self):
        ret_c, system_part = self.cmd("df -h / | grep -v Filesystem | awk '{print $1}'")
        ret_c, os_disk = self.cmd("lsblk | grep disk | awk '{print $1}' | xargs")
        for disk in os_disk.split():
            if disk in system_part:
                continue
            else:
                ret_c,ret_o = self.cmd("lsblk | grep %s | awk '{print $4}'" % disk)
                if "0B" == ret_o:
                    continue
                else:
                    return disk
